package com.example.deepai.util;

import io.milvus.client.MilvusServiceClient;
import io.milvus.param.dml.QueryParam;
import io.milvus.param.collection.ShowCollectionsParam;
import io.milvus.param.collection.GetCollectionStatisticsParam;
import io.milvus.param.R;
import io.milvus.grpc.CheckHealthResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import java.util.Arrays;

@Slf4j
@Component
public class MilvusQueryUtil {
    private final MilvusServiceClient milvusClient;
    
    public MilvusQueryUtil(MilvusServiceClient milvusClient) {
        this.milvusClient = milvusClient;
    }
    
    // Fix health check response type
    public boolean checkMilvusStatus() {
        try {
            R<CheckHealthResponse> healthResponse = milvusClient.checkHealth();
            return healthResponse != null && healthResponse.getData() != null;
        } catch (Exception e) {
            log.error("Milvus 服务状态检查失败", e);
            return false;
        }
    }
    
    // 列出所有集合
    public void listCollections() {
        try {
            var result = milvusClient.showCollections(ShowCollectionsParam.newBuilder().build());
            log.info("现有集合：{}", result.getData().getCollectionNamesList());
        } catch (Exception e) {
            log.error("获取集合列表失败", e);
        }
    }
    
    // 查询集合中的数据
    // Fix query limit type
    public void queryCollection(String collectionName) {
        try {
            QueryParam queryParam = QueryParam.newBuilder()
                .withCollectionName(collectionName)
                .withOutFields(Arrays.asList("id", "feature_vector"))
                .withLimit(10L)  // Change to Long type
                .build();
            
            var result = milvusClient.query(queryParam);
            log.info("查询结果：{}", result.getData());
        } catch (Exception e) {
            log.error("查询集合数据失败", e);
        }
    }
    
    // 获取集合统计信息
    public void getCollectionStats(String collectionName) {
        try {
            var stats = milvusClient.getCollectionStatistics(
                GetCollectionStatisticsParam.newBuilder()
                    .withCollectionName(collectionName)
                    .build()
            );
            log.info("集合统计信息：{}", stats.getData());
        } catch (Exception e) {
            log.error("获取集合统计失败", e);
        }
    }
}